from mindspore import Tensor
import mindspore as ms
import mindspore.numpy as np
eps=1e-8
EPS = Tensor(float(0.00000001), dtype=ms.float32)

def cross(a, b):
	return a[0] * b[1] - a[1] * b[0]

def point_cmp(a, b, center):
	atan2=ms.ops.Atan2()
	a = atan2(a[1] - center[1], a[0] - center[0]) > atan2(b[1] - center[1], b[0] - center[0])
	return a

def cross_point(p1, p2, p0):
	return (p1[0] - p0[0]) * (p2[1] - p0[1]) - (p2[0] - p0[0]) * (p1[1] - p0[1])

def check_rect_cross(p1, p2, q1, q2):
	ret = (min(p1[0], p2[0]) <= max(q1[0], q2[0])) and (min(q1[0], q2[0]) <= max(p1[0], p2[0])) and \
		  (min(p1[1], p2[1]) <= max(q1[1], q2[1])) and (min(q1[1], q2[1]) <= max(p1[1], p2[1]))
	return ret

def rotate_around_center(center, angle_cos, angle_sin, p):
	new_x = (p[0] - center[0]) * angle_cos + (p[1] - center[1]) * (-angle_sin) + center[0]
	new_y = (p[0] - center[0]) * angle_sin + (p[1] - center[1]) * angle_cos + center[1]
	p[0] = new_x
	p[1] = new_y


def check_in_box2d(box, p):
	# //params: (7) [x, y, z, dx, dy, dz, heading]
	MARGIN = Tensor(float(1e-2),dtype=ms.float32)

	center_x = box[0]
	center_y = box[1]
	
	angle_cos = np.cos(-box[6])
	angle_sin = np.sin(-box[6])  # // rotate the point in the opposite direction of box
	rot_x = (p[0] - center_x) * angle_cos + (p[1] - center_y) * (-angle_sin)
	rot_y = (p[0] - center_x) * angle_sin + (p[1] - center_y) * angle_cos
	
	return (rot_x.abs() < box[3] / 2 + MARGIN) and (rot_y.abs() < box[4] / 2 + MARGIN)

def intersection(p1, p0, q1, q0, ans):
	# // fast exclusion
	if check_rect_cross(p0, p1, q0, q1) == 0:
		return 0

	# // check cross standing
	s1 = cross_point(q0, p1, p0)
	s2 = cross_point(p1, q1, p0)
	s3 = cross_point(p0, q1, q0)
	s4 = cross_point(q1, p1, q0)

	if not (s1 * s2 > 0 and s3 * s4 > 0):
		return 0

	# // calculate intersection of two lines
	s5 = cross_point(q1, p1, p0)
	# if ops.abs(s5 - s1) > EPS:
	if (s5 - s1).abs() > EPS:
		ans[0] = (s5 * q0[0] - s1 * q1[0]) / (s5 - s1)
		ans[1] = (s5 * q0[1] - s1 * q1[1]) / (s5 - s1)
	else :
		a0 = p0[1] - p1[1]
		b0 = p1[0] - p0[0]
		c0 = p0[0] * p1[1] - p1[0] * p0[1]
		a1 = q0[1] - q1[1]
		b1 = q1[0] - q0[0]
		c1 = q0[0] * q1[1] - q1[0] * q0[1]
		D = a0 * b1 - a1 * b0
		ans[0] = (b0 * c1 - b1 * c0) / D
		ans[1] = (a1 * c0 - a0 * c1) / D
	return 1


def box_overlap(box_a, box_b):
	# params box_a: [x, y, z, dx, dy, dz, heading]
	# params box_b: [x, y, z, dx, dy, dz, heading]
	a_angle = box_a[6]
	b_angle = box_b[6]
	a_dx_half = box_a[3] / 2
	b_dx_half = box_b[3] / 2
	a_dy_half = box_a[4] / 2
	b_dy_half = box_b[4] / 2
	a_x1 = box_a[0] - a_dx_half
	a_y1 = box_a[1] - a_dy_half
	a_x2 = box_a[0] + a_dx_half
	a_y2 = box_a[1] + a_dy_half
	b_x1 = box_b[0] - b_dx_half
	b_y1 = box_b[1] - b_dy_half
	b_x2 = box_b[0] + b_dx_half
	b_y2 = box_b[1] + b_dy_half

	center_a = np.append(box_a[0], box_a[1]) 
	center_b = np.append(box_b[0], box_b[1]) 
	
	box_a_corners = np.concatenate((np.append(a_x1, a_y1), np.append(a_x2, a_y1), np.append(a_x2, a_y2), np.append(a_x1, a_y2), np.array([0, 0])),axis=0).reshape(5,2)
	box_b_corners = np.concatenate((np.append(b_x1, b_y1), np.append(b_x2, b_y1), np.append(b_x2, b_y2), np.append(b_x1, b_y2), np.array([0, 0])), axis=0).reshape(5,2)

	a_angle_cos = np.cos(a_angle)
	a_angle_sin = np.sin(a_angle)
	b_angle_cos = np.cos(b_angle)
	b_angle_sin = np.sin(b_angle)

	for k in range(4):
		rotate_around_center(center_a, a_angle_cos, a_angle_sin, box_a_corners[k, :])
		rotate_around_center(center_b, b_angle_cos, b_angle_sin, box_b_corners[k, :])

	box_a_corners[4, :] = box_a_corners[0, :]
	box_b_corners[4, :] = box_b_corners[0, :]

	# // get intersection of lines
	cross_points = Tensor([[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],[0,0],
								[0,0],[0,0]], dtype=ms.float32)
	cnt = 0
	flag = 0
	poly_center = Tensor([0, 0], dtype=ms.float32)
	for i in range(4):
		for j in range(4):
			flag = intersection(box_a_corners[i + 1, :], box_a_corners[i, :], box_b_corners[j + 1, :],
								box_b_corners[j, :], cross_points[cnt, :])
			if flag:
				poly_center = poly_center + cross_points[cnt, :]
				cnt += 1

	# // check corners
	for k in range(4):
		if check_in_box2d(box_a, box_b_corners[k,:]):
			poly_center[0] = poly_center[0] + box_b_corners[k, 0]
			poly_center[1] = poly_center[1] + box_b_corners[k, 1]
			cross_points[cnt,:] = box_b_corners[k, :]
			cnt += 1
		if check_in_box2d(box_b, box_a_corners[k,:]):
			poly_center[0] = poly_center[0] + box_b_corners[k, 0]
			poly_center[1] = poly_center[1] + box_b_corners[k, 1]
			cross_points[cnt,:] = box_a_corners[k,:]
			cnt +=1

	poly_center[0] /= Tensor(cnt,ms.int32)
	poly_center[1] /= Tensor(cnt,ms.int32)

	# // sort the points of polygon
	for j in range(cnt - 1):
		for i in range(cnt-j - 1):
			if point_cmp(cross_points[i, :], cross_points[i + 1, :], poly_center):
				temp = cross_points[i, :].copy()
				cross_points[i,:] = cross_points[i + 1,:]
				cross_points[i + 1,:] = temp


	# // get the overlap areas
	area = Tensor(0.0, dtype=ms.float32)
	cps1 = Tensor([0, 0], dtype=ms.float32)
	cps2 = Tensor([0, 0], dtype=ms.float32)
	for k in range(cnt-1):
		cps1[0] = cross_points[k, 0] - cross_points[0, 0]
		cps1[1] = cross_points[k, 1] - cross_points[0, 1]
		cps2[0] = cross_points[k+1, 0] - cross_points[0, 0]
		cps2[1] = cross_points[k+1, 1] - cross_points[0, 1]
		area += cross(cps1, cps2)
	return area.abs() / 2.0

def iou_bev(box_a, box_b):
	# params
	# box_a: [x, y, z, dx, dy, dz, heading]
	# params
	# box_b: [x, y, z, dx, dy, dz, heading]
	# 计算两个box的面积
	sa = box_a[3] * box_a[4]
	sb = box_b[3] * box_b[4]
	# 计算交集面积
	s_overlap = box_overlap(box_a, box_b)
	#  iou计算
	return np.divide(s_overlap.astype(ms.float32), np.maximum(sa + sb - s_overlap, EPS).astype(ms.float32))


def nms_gpu(boxes, scores, thresh, pre_maxsize=None, **kwargs):
	"""
	:param boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
	:param scores: (N)
	:param thresh:
	:return:
	"""
	assert boxes.shape[1] == 7
	# 对分数按列降序排序(从大到小)，并取出对应索引
	# dim=0 按列排序，dim=1 按行排序，默认 dim=1
	# 因为传入的scores已经有序，所以order就是[1 2 3 ...]
	sort=ms.ops.Sort(axis=0, descending=True)
	order = sort(scores.astype(ms.float16))[1].astype(ms.int32)
	# 如果存在NMS前的最大box数量（4096），则取出前4096个box索引
	# print(order)
	if pre_maxsize is not None:
		order = order[:pre_maxsize]

	# keep = np.array([0],dtype=ms.int32)
	keep=[]
	# keep：记录保留目标框的下标
	# num_out：返回保留下来的个数
	size=ms.ops.Size()
	while np.size(order) > 0:
		if np.size(order) == 1:
			# i = order.item()
			keep.append(int(order))
			break
		else:
			keep.append(int(order[0]))
		iou=np.array([0.1],dtype=ms.float32)
		for box in boxes[order[1:]]:
			iou = np.append(iou,iou_bev(boxes[order[0]], box))
		iou=iou[1:]
		idx = np.squeeze(ms.ops.nonzero(iou <= thresh)).astype(ms.int32)
		order=order.take(Tensor(idx.asnumpy()+1))

		
	return Tensor(keep)

if __name__ == "__main__":
	# print(nms_gpu(torch.tensor([2,4,5,0.2,0.3,0.4,2]).reshape((1,7)), torch.tensor([0.2]).reshape(1),0.8))
	# import numpy
	# a = Tensor(np.array([4, 3, 5, 7, 6, 8]))
	# indices = Tensor(np.array([0, 1, 4]))
	# output = a.take(indices)
	# print(output)
	# 32401
	print(nms_gpu(Tensor([[-8.545734 ,-30.78366 ,    -2.0139227 ,   0.5226521  ,  2.7160504,
    0.92416924 , -2.7668417 ],
 [ -9.075958  , -30.09282    , -1.9771283  ,  0.5064265  ,  2.2899377,
    0.92757404  , 2.656907  ],
 [ -2.9210777,   21.15643   ,  -0.70781124 ,  0.57447064 ,  2.5526633,
    0.95974404 , -3.1383827 ],
 [ 18.89481 ,    49.26737   ,  -0.585745   ,  0.6211397   , 3.0566106,
    0.9712022  , -1.100473  ],
 [ 19.49757   ,  49.124363  ,  -0.7543063   , 0.61432874  , 2.9472034,
    0.9668016 ,  -1.0049689 ]]).reshape((5,7)), Tensor([0.11718041, 0.10943048 ,0.12485003, 0.13310847, 0.119372 ]).reshape(5),Tensor(0.2).reshape(1)))